# -*- coding: utf-8 -*-

from StoreBilibiliData.CreateTable import metadata, uploader, video, comments, dm

from sqlalchemy import create_engine, insert, select, and_
from sqlalchemy.sql import func
from sqlalchemy.sql.dml import Insert
from sqlalchemy.sql.selectable import Select
from sqlalchemy.engine.base import Engine, Connection
from sqlalchemy.engine.result import ResultProxy, RowProxy

import os
import re
import datetime

'''
按月合并数据库文件；
'''


def database_migration(db_new: str, *db_old: str) -> None:
    engine_new = create_engine('sqlite:///' + db_new, echo=True, encoding='utf-8')
    engine_mem = create_engine('sqlite:///:memory:', echo=True, encoding='utf-8')
    con_new = engine_new.connect()  # type: Connection
    con_mem = engine_mem.connect()  # type: Connection
    metadata.create_all(engine_new)
    metadata.create_all(engine_mem)

    con_old = []
    for db in db_old:
        engine_old = create_engine('sqlite:///' + db, echo=True, encoding='utf-8')
        con_old.append(engine_old.connect())

    def __merge_uploader() -> None:
        sel_old = select([uploader])
        sel_new = select([uploader.c.uid])
        ins_new = insert(uploader)

        for con in con_old:
            rp_old = con.execute(sel_old)
            for up in rp_old:
                tmp_rp = con_mem.execute(sel_new.where(uploader.c.uid == up.uid))
                if not tmp_rp.first():
                    con_mem.execute(ins_new.values(
                        uid=up.uid,
                        name=up.name,
                    ))

    def __merge_video() -> None:
        sel_old = select([video])  # type: Select
        sel_new = select([video.c.av])  # type: Select
        ins_new = insert(video)  # type: Insert

        for con in con_old:  # type: Connection
            rp_old = con.execute(sel_old)  # type: ResultProxy
            for v in rp_old:  # type: RowProxy
                tmp_rp = con_mem.execute(sel_new.where(video.c.av == v.av))  # type: ResultProxy
                if not tmp_rp.first():
                    con_mem.execute(ins_new.values(
                        av=v.av,
                        bv=v.bv,
                        comment_count=v.comment_count,
                        play_count=v.play_count,
                        title=v.title,
                        description=v.description,
                        uploader_id=v.uploader_id,
                        upload_time=v.upload_time,
                    ))

    def __merge_comments() -> None:
        sel_old = select([comments])
        sel_new = select([comments.c.rp_id])
        ins_new = insert(comments)

        for con in con_old:
            rp_old = con.execute(sel_old)
            for c in rp_old:
                tmp_rp = con_mem.execute(sel_new.where(comments.c.rp_id == c.rp_id))  # type: ResultProxy
                if not tmp_rp.first():
                    con_mem.execute(ins_new.values(
                        rp_id=c.rp_id,
                        video_id=c.video_id,
                        likes=c.likes,
                        root_comment=c.root_comment,
                        content=c.content,
                        comment_time=c.comment_time,
                    ))

    def __merge_dm() -> None:
        sel_old = select([dm])
        sel_new = select([dm.c.id])
        ins_new = insert(dm)

        for con in con_old:
            rp_old = con.execute(sel_old)
            for d in rp_old:
                tmp_rp = con_mem.execute(sel_new.where(and_(dm.c.content == d.content, dm.c.property == d.property)))
                if not tmp_rp.first():
                    con_mem.execute(ins_new.values(
                        video_id=d.video_id,
                        content=d.content,
                        property=d.property,
                    ))

    def __from_mem_to_disk() -> None:
        uploader_rp = con_mem.execute(select([uploader]))
        video_rp = con_mem.execute(select([video]))
        comments_rp = con_mem.execute(select([comments]))
        dm_rp = con_mem.execute(select([dm]))

        uploader_dict = []
        video_dict = []
        comments_dict = []
        dm_dict = []

        for r in uploader_rp:
            uploader_dict.append({
                'uid': r.uid,
                'name': r.name,
            })

        for r in video_rp:
            video_dict.append({
                'av': r.av,
                'bv': r.bv,
                'comment_count': r.comment_count,
                'play_count': r.play_count,
                'title': r.title,
                'description': r.description,
                'uploader_id': r.uploader_id,
                'upload_time': r.upload_time,
            })

        for r in comments_rp:
            comments_dict.append({
                'rp_id': r.rp_id,
                'video_id': r.video_id,
                'likes': r.likes,
                'root_comment': r.root_comment,
                'content': r.content,
                'comment_time': r.comment_time,
            })

        for r in dm_rp:
            dm_dict.append({
                'video_id': r.video_id,
                'content': r.content,
                'property': r.property,
            })

        con_new.execute(
            insert(uploader),
            uploader_dict,
        )

        con_new.execute(
            insert(video),
            video_dict,
        )

        con_new.execute(
            insert(comments),
            comments_dict,
        )

        con_new.execute(
            insert(dm),
            dm_dict,
        )

    __merge_uploader()
    __merge_video()
    __merge_comments()
    __merge_dm()
    __from_mem_to_disk()


def check_migration_result(db_new: str, *db_old: str) -> str:
    con_new = create_engine("sqlite:///" + db_new, echo=True, encoding='utf-8').connect()  # type: Connection
    con_old = []
    for db in db_old:
        con_old.append(
            create_engine('sqlite:///' + db, echo=True, encoding='utf-8').connect()
        )

    video_count_new = 0
    video_count_old = 0

    comment_count_new = 0
    comment_count_old = 0

    dm_count_new = 0
    dm_count_old = 0

    for con in con_old:
        sel = select([func.count(video.c.av).label('cnt')])
        rp = con.execute(sel)  # type: ResultProxy
        video_count_old += rp.first().cnt

        sel = select([func.count(comments.c.rp_id).label('cnt')])
        rp = con.execute(sel)  # type: ResultProxy
        comment_count_old += rp.first().cnt

        sel = select([func.count(dm.c.id).label('cnt')])
        rp = con.execute(sel)  # type: ResultProxy
        dm_count_old += rp.first().cnt

        if video_count_new == 0:
            sel = select([func.count(video.c.av).label('cnt')])
            rp = con_new.execute(sel)  # type: ResultProxy
            video_count_new = rp.first().cnt

        if comment_count_new == 0:
            sel = select([func.count(comments.c.rp_id).label('cnt')])
            rp = con_new.execute(sel)  # type: ResultProxy
            comment_count_new = rp.first().cnt

        if dm_count_new == 0:
            sel = select([func.count(dm.c.id).label('cnt')])
            rp = con_new.execute(sel)  # type: ResultProxy
            dm_count_new = rp.first().cnt

    return f'{video_count_new=}\n{video_count_old=}\n\n' + \
           f'{comment_count_new=}\n{comment_count_old=}\n\n' + \
           f'{dm_count_new=}\n{dm_count_old=}\n'


if __name__ == '__main__':
    program_start = datetime.datetime.now()
    data_folder = 'D:/MyResources/Web Scraping/Bilibili/'
    merged_db_name = './DatabaseFile/bilibili 2020-04.db'
    files = os.listdir(data_folder)
    db_to_merge = []
    matcher = re.compile('bilibili 2020-04-[0-9]{2} - 2020-04-[0-9]{2}.db')
    for f in files:
        if matcher.match(f):
            db_to_merge.append(data_folder + f)

    database_migration(merged_db_name, *db_to_merge)
    res = check_migration_result(merged_db_name, *db_to_merge)
    print(res)
    program_end = datetime.datetime.now()
    print(program_start)
    print(program_end)
    print(program_end - program_start)
